#!/usr/bin/env python
"""
:mod:`pynchon.bio.alg.seq`
==========================

Algorithms on sequences.
"""
import numpy as np
import scipy as sp
from itertools import izip, combinations, permutations
from pynchon.core.index import IndexVector0dInt8
from pynchon.core.point import PointVector0dFloat64
from pynchon.math.array import max_index
from pynchon.bio.data.score import BLOSUM62
from pynchon.bio.aa import AASeq
from pynchon.bio.nt import Nt

def score_seqs(seq1, seq2, sm=BLOSUM62):
    """
    """
    score = []
    if len(seq1) != len(seq2):
        raise ValueError('Sequences have different length.')
    for s1,s2 in izip(seq1, seq2):
        score.append(sm[s1][s2])
    return IndexVector0dInt8(score)
         
def hamming_dist(seq1, seq2):
    """
    Returns the Hamming distance between two sequences. The Hamming distance is 
    the number of characters which differ between two sequences (arrays) of 
    equal length.
    
    Arguments:
    
        - seq1 (``numpy.ndarray``) ``SymbolVector`` or rank-1 array of symbols.
        - se12 (``numpy.ndarray``) ``SymbolVector`` or rank-1 array of symbols.
    """
    if not (len(seq1.shape) == len(seq2.shape) == 1):
        raise ValueError("Arrays have to be of rank 1.")
    if not (seq1.shape[0] == seq2.shape[0]):
        raise ValueError("Arrays have to be of equal length.")  
    return np.sum(seq1 != seq2, axis=-1) 

def all_pairs(seq, symmetric =False):
    """
    Returns all pairs of symbols within a sequence.
    
    Arguments:
    
        - seq (``sequence`` or ``SymbolVecotr``) Any sequence of symbols or a 
          ``SymbolVector``.
        - symmetric (``bool``) [default: False] If ``True`` only one off 
          diagonal pair is returned.
    """
    if symmetric:
        pairs = list(combinations(seq, r=2))
    else:
        pairs = list(permutations(seq, r=2))
    return pairs

def sum_of_pairs(seq, normalized =False, sm=BLOSUM62):
    """
    Calculates the sum-of-pairs score. Scores are calculated for sequences along
    the last axis in the seq array. Assumes that the score matrix is symmetric
    and so a only one pair of (posN,posM) and (posM,posN) is counted. 
    
    Arguments:
    
        - seq (``numpy.ndarray``) A ``Symbol`` instance of at least rank-1 or
          ``numpy.ndarray``.
        - normalized (``bool``) If ``True`` the score will be normalized by the
          number of pairs i.e. by 2 / (N * (N - 1))
        - sm (``dict``) [default: BLOSUM62] The score matrix to be used.
    """
    shape = seq.shape # keep shape to return in right shape
    scores = np.zeros(seq.size / shape[-1]) # place to keep scores
    seq = seq.reshape((-1, shape[-1])) 
    seq = np.atleast_2d(seq)
    for i, s in enumerate(seq):
        pairs = 0
        for sym1, sym2 in all_pairs(s, symmetric=True):
            try:
                scores[i] += sm[sym1][sym2]
                pairs += 1
            except KeyError:
                pass
        if normalized:
            scores[i] *= (2. /(pairs * (pairs - 1)))
    scores.reshape(shape[:-1])
    return PointVector0dFloat64(scores)
   
def align(seqj, seqi, gap_open=-5, gap_extend=-1, local=False, sm=BLOSUM62):
    """
    Calculates the alignment of two sequences using the Needleman-Wunsch (local
    ``False``) or Smith-Waterman (local ``True``) algorithm. Returns the aligned
    sequences as character arrays.
    
    Arguments:
    
        - seqj (``sequence``) First aligned iterable object of symbols.
        - seqi (``sequence``) Second aligned iterable object of symbols.
        - gap_open (``float``) The gap-opening cost.
        - gap_extend (``float``) The gap-extension cost.
        - local (``bool``) If ``True`` will return the best local alignment 
          between the sequences.
    """
    NONE, LEFT, UP, DIAG = range(4) # NONE is 0
    max_j = len(seqj)
    max_i = len(seqi)
    
    if max_j > max_i:
        flip = 1
        seqi, seqj = seqj, seqi
        max_i, max_j = max_j, max_i
    else:
        flip = 0
    
    score   = np.zeros((max_i + 1, max_j + 1), dtype=np.float32)
    pointer = np.zeros((max_i + 1, max_j + 1), dtype=np.uint) # NONE
    agap_i = np.zeros((max_i + 1,), dtype=np.bool)
    agap_j = np.zeros((max_j + 1,), dtype=np.bool)
    
    if not local:
        pointer[0, 1:] = LEFT
        pointer[1:, 0] = UP
        score[0, 1:] = gap_open + gap_extend * np.arange(0, max_j, dtype=np.float32)
        score[1:, 0] = gap_open + gap_extend * np.arange(0, max_i, dtype=np.float32)
    
    for i in range(1, max_i + 1):
        ci = seqi[i - 1]
        for j in range(1, max_j + 1):
            cj = seqj[j - 1]
            
            diag_score = score[i - 1, j - 1] + sm[cj][ci]
            left_score = score[i, j - 1] + (gap_open if agap_j[j - 1] == False else gap_extend)
            up_score   = score[i - 1, j] + (gap_open if agap_i[j] == False else gap_extend)
            max_score = max(diag_score, up_score, left_score)
            
            score[i,j] = max(0, max_score) if local else max_score

            if local and score[i,j] == 0:
                # leave pointer NONE
                agap_j[j] = False
                agap_i[j] = False                 
            elif max_score == diag_score:
                pointer[i,j] = DIAG
                agap_j[j] = False
                agap_i[j] = False    
            elif max_score == left_score:
                pointer[i,j] = LEFT
                agap_j[j] = True
                agap_i[j] = False
            else:
                pointer[i,j] = UP
                agap_j[j] = False
                agap_i[j] = True        
    align_j = []
    align_i = []
    if local:
        i,j = max_index(score)
    p = pointer[i, j] 
    while p != NONE:
        if p == DIAG:
            i -= 1
            j -= 1            
            align_j.append(seqj[j])
            align_i.append(seqi[i])
        elif p == LEFT:
            j -= 1            
            align_j.append(seqj[j])
            align_i.append("-")
        elif p == UP:
            i -= 1            
            align_j.append("-")
            align_i.append(seqi[i])
        else:
            raise Exception('wtf!')
        p = pointer[i, j]
    np.array(align_j.reverse())
    np.array(align_i.reverse())
    return ((align_i, align_j) if flip else (align_j, align_i))


if __name__ == '__main__':
    if 1:
        a,b = align('A', 'A', -7, -2)
        assert a == ['A']
        assert b == ['A']
        a,b = align('R','K', -7, -2)
        assert a == ['R']
        assert b == ['K'] 
        a,b = align('R','AR', -7, -2)   
        assert a == ['-','R']
        assert b == ['A','R']   
        a,b = align('AR','R', -7, -2)   
        assert a == ['A','R']
        assert b == ['-','R']  
        a,b = align('AR','RA', -7, -2)   
        assert a == ['A','R']
        assert b == ['R','A']    
        a,b = align('AR','RA', -3, -2)      
        assert a == ['A', 'R', '-']
        assert b == ['-', 'R', 'A']      
        a,b = align('RAR','RR', -3, -2)
        assert b == ['R', '-', 'R']
        assert a == ['R', 'A', 'R']    
        a,b = align('RAR','RR', -10, -2)
        assert b == ['R', '-', 'R']
        assert a == ['R', 'A', 'R']      
        a,b = align('RAAR','RR', -9, -5)
        assert a == ['R', 'A', 'A', 'R']
        assert b == ['R', '-', '-', 'R']
        a,b = align('RLR','RER', -9, -5)
        assert a == ['R', 'L', 'R']
        assert b == ['R', 'E', 'R'] 
        a,b = align('RLR','RER', -1, -.5)
        assert a == ['R', '-', 'L', 'R']
        assert b == ['R', 'E', '-', 'R']
        a,b = align('RLR','REER', -1, -.5)
        assert a == ['R', '-', '-', 'L', 'R']
        assert b == ['R', 'E', 'E', '-', 'R']     
        a, b = align('AGEBAM', 'AGEBAMAM', -6, -2)   
        assert a == ['A', 'G', 'E', 'B', '-', '-', 'A', 'M']
        assert b == ['A', 'G', 'E', 'B', 'A', 'M', 'A', 'M']
        a, b= align('CPELIRKNCANTH', 'PREKRLICAN', -0.5, -0.1)
        assert list(a) == ['C', 'P', '-', 'E', '-', '-', 'L', 'I', 'R', 'K', 'N', 'C', 'A', 'N', 'T', 'H']
        assert list(b) == ['-', 'P', 'R', 'E', 'K', 'R', 'L', 'I', '-', '-', '-', 'C', 'A', 'N', '-', '-']    
        a, b= align('CPEL', 'PREK', -6., -2.)
        assert a == ['C', 'P', 'E', 'L']
        assert b == ['P', 'R', 'E', 'K']  
        a, b= align('CPEL', 'PREK', -5, -2.)
        assert a == ['C', 'P', '-','E', 'L']
        assert b == ['-','P', 'R', 'E', 'K']    
        a,b = align('RLRR','RRER', -1, -.5)
        assert a == ['R', 'L', 'R', '-', 'R']
        assert b == ['R', '-', 'R', 'E', 'R']
    
    a, b = align('A', 'A', -7., -2., True)
    assert a == ['A']
    assert b == ['A']

    a, b = align('RA', 'AR', -7., -2., True)
    assert a == ['R']
    assert b == ['R']    
    a, b = align('RRR', 'RR', -7., -2., True)
    assert a == ['R', 'R']
    assert b == ['R', 'R']
    a,b = align('PYNCHAN', 'YNCH', -7., -2., True)
    assert a == ['Y', 'N', 'C', 'H']
    assert b == ['Y', 'N', 'C', 'H']        
    a, b = align('AIP', 'AP', -6   , -2., True)
    assert a == ['P']
    assert b == ['P']   
    
    a, b = align('AIP', 'AP', -2.   , -1., True)
    ['A', 'I', 'P']
    ['A', '-', 'P']
    assert a == a
    assert b == b   
    a, b = align('RYY', 'RAA', -7., -2., True)
    assert a == ['R']
    assert b == ['R']   
    a, b = align('RYY', 'AAR', -7., -2., True)
    assert a == ['R']
    assert b == ['R']        
    a, b = align('R', 'R', -7., -2., True)
    assert a == ['R']
    assert b == ['R']
    a, b = align('REE', 'LLR', -7., -2., True)
    assert a == ['R']
    assert b == ['R']    
    a, b = align('RRR', 'AAA', -7., -2., True)
    print a, b
    assert a == ['R']
    assert b == ['A']   
   

    #needle asis:CPELINCANTH asis:PRELICAN -gapopen=1 -gapextend=1 -stdout -auto
#    a,b = global_align('CPEL', 'PREK', -7, -2)
#    print a
#    print b
#    assert b == ['P','R','E','K']
#    a, b = global_align('AGEBAM', 'AGEBAMAM', -6, -2)

#    a, b = global_align('WIPP', 'WPP', -6, -2)
#    assert a == ['W','I','P','P']
#    assert b == ['W','-','P','P']
#    a, b = global_align('WRAGI', 'WIRGP', -6, -2)
#    assert a == ['W','R','A','G','I']
#    assert b == ['W','I','R','G','P']          




#    a, b = local_align('AP', 'PA', -7   , -2.)
#    assert a == ['P']
#    assert b == ['P']
#    a, b = local_align('PA', 'PA', -7   , -2.)
#    assert a == ['P', 'A']
#    assert b == ['P', 'A']    
#    a, b = local_align('PIA', 'PA', -6   , -2.)
#    assert a == ['P']
#    assert b == ['P']         
  
#    

#    #assert a == ['P']
#    #assert b == ['P']     
#    print a 
#    print b
#    
    
    
    